#!/bin/env python

###### Functions for plotting, Camille Avestruz, Last updated 5/21/2012 ######
###### Note: only functions with ## have been checked, else in progress ######

from re import search
from glob import glob
from numpy import array, sqrt
import sys

def find_nearest_profiles_aexpn( profiledir, aexpn = 1.0 ) : ## Takes in a profiles directory, and a float. Returns the string of the nearest aexpn
    nearest_aexp = '-1'
    for output in glob( profiledir + '/h_blist_a*.dat') :
        aexp = search('h_blist_a(\d\.\d\d\d\d).dat',output).group(1)
        if ( abs(float(aexp)-aexpn) < abs(float(nearest_aexp)-aexpn) ) :
            nearest_aexp = aexp
    return nearest_aexp # String

def blist_fractions( blist, ids, OmbOmMratio = 0.143 ) : ## This takes in a loaded h_blist and outputs Mvir, fgas, fstar, fbary for a list of halo ids in terms of Omega baryon to Omega matter ratio

    Mvir_dict = dict(zip(ids,[blist[7][id] for id in ids ]))
    fgas_dict = dict(zip(ids,[(blist[1][id]-blist[2][id])/blist[7][id]/OmbOmMratio for id in ids]))
    fgastot_dict = dict(zip(ids,[blist[1][id]/blist[7][id]/OmbOmMratio for id in ids]))
    fstar_dict = dict(zip(ids,[blist[3][id]/blist[7][id]/OmbOmMratio for id in ids]))
    fbary_dict = dict(zip(ids,[blist[5][id]/blist[7][id]/OmbOmMratio for id in ids]))

    blist_fractions = { 'Mvir':Mvir_dict, 'fgas':fgas_dict, 'fgastot':fgastot_dict, 'fstar':fstar_dict, 'fbary':fbary_dict }

    return blist_fractions # Access as blist_fractions['Mvir'][id], etc.

def halo_order( blist ) : # Returns a list of halo_ids in order of either decreasing vmax or decreasing halo mass

    vmax_pairs = zip(blist[8],blist[8].keys())
    Mhalo_pairs = zip(blist[7],blist[7].keys())

    sorted_ids_vmax = [ pair[1] for pair in sorted(vmax_pairs, key=lambda row: row[0], reverse=True ) ]
    sorted_ids_Mhalo = [ pair[1] for pair in sorted(Mhalo_pairs, key=lambda row: row[0], reverse=True ) ]

    ordered_halo_ids = { 'vmax':sorted_ids_vmax, 'Mhalo':sorted_ids_Mhalo }

    return ordered_halo_ids # Access as ordered_halo_ids['vmax'][0] for the halo_id with largest vmax, ordered_halo_ids['Mhalo'][0] for the halo_id with the largest Mhalo

def read_print_agn( workdir, aexpn = 1.0 ) : ## Returns dictionary of agn masses and position arrays with id as key - all the AGN_* functions use this 
    printagndir = workdir + '/logs/print_agn/'
    profiledir = workdir + '/profiles/'
    aexp = find_nearest_profiles_aexpn( profiledir, aexpn )
    f = open(printagndir+'print_agn_a'+aexp+'.dat','r') # id Mbh x y z
    agn_lines = [ line.split() for line in f ]
    printed_agn = dict(zip([agn_line[0] for agn_line in agn_lines],
                           [{ 'Mbh':float(agn_line[1]),
                              'BHx':array( [float(agn_line[i]) for i in range(2,5)] ) }
                            for agn_line in agn_lines
                             ]))
    f.close()
    
    return printed_agn # Access as printed_agn[id]['Mbh'] or printed_agn[id]['BHx'] to get position array

def AGN_mass_sort( workdir, aexpn = 1.0 ) : ## Returns the particle_ids of all of the AGN in order of decreasing mass at aexpn 
    printed_agn = read_print_agn( workdir, aexpn )
    printed_agn_by_mass = [ (printed_agn[id]['Mbh'], id) for id in printed_agn.keys() ]

    mass_sorted_ids = [ mass_sorted_id[1] for mass_sorted_id in sorted( printed_agn_by_mass, key=lambda row: row[0], reverse=True ) ]

    return mass_sorted_ids 

def AGN_in_sphere( workdir, sphere_center, ngrid, Lbox, radius, aexpn ) : ## Takes in work directory, sphere center array in code units, float ngrid (128), float Lbox, and radius in comoving Mpc/h.  Returns AGN ids within that sphere
    radius_in_code_units = (ngrid/Lbox) * radius
    printed_agn = read_print_agn( workdir, aexpn )
    printed_agn_in_radius_ids = [ id for id in printed_agn.keys() if sqrt(sum((printed_agn[id]['BHx']-sphere_center)**2)) < radius_in_code_units ]
    
    return printed_agn_in_radius_ids

def AGN_total_mass( workdir, aexpn = 1.0 ) : ## Returns the total black hole mass at aexpn 
    printed_agn = read_print_agn( workdir, aexpn )
    
    return sum( printed_agn[id]['Mbh'] for id in printed_agn.keys() )

def AGN_progenitors( workdir, particle_id ) : ## Returns particle_ids of the progenitors of particle_id as keys, with values aexpn (of merger) and Mbh 
    logdir = workdir+'/logs/'
    progenitors_of_agn = {}
    for mergersfile in glob(logdir+'Merged_BlackHoles*.log') :
        f = open(mergersfile, 'r')
        for line in f :
            if not line.startswith('#') :
                cols = line.split()
                if cols[1] == particle_id :
                    ( merged_id, merged_aexp, merged_mass) = ( cols[16], float(cols[0]), float(cols[17]) )
                    progenitors_of_agn[merged_id] = { 'merged_mass':merged_mass,'merged_aexp':merged_aexp }
                    
    return progenitors_of_agn # Access as AGN_progenitors[progenitor_id]['merged_aexp'] 
    
# def AGN_host_halo( workdir, r, aexp = 1.0 ) : # For each AGN, find the halo_id it is closest to and ratio of r ('r200', 'r500', 'rvir', etc.) it is from the center. 
    ###
